import { Algorithm } from "../types"

// 数据模型定义
interface MPLPNode {
  id: any
  type: string
  attributes: Record<string, any>
}

interface MPLPEdge {
  source: any
  target: any
  type?: string
  weight?: number
  attributes: Record<string, any>
}

interface MPLPGraph {
  nodes: MPLPNode[]
  edges: MPLPEdge[]
}

// 定义时序图数据结构
interface TemporalGraph {
  nodes: MPLPNode[]
  edges: MPLPEdge[]
  timestamp: string
}

// 修改任务状态枚举
enum TaskStatus {
  PENDING = "PENDING",
  RUNNING = "RUNNING", 
  COMPLETED = "COMPLETED",
  FAILED = "FAILED"
}

// 修改接口定义以匹配后端schema
interface TSGEInputParams {
  graph_sequence: Array<{
    nodes: MPLPNode[];
    edges: MPLPEdge[];
  }>;
  embedding_dim: number;
  learning_rate: number;
  epochs: number;
  node_types?: Record<string, string>;
}

interface NodeEmbedding {
  node_id: string;
  embedding: number[];
  node_type?: string;
}

interface LinkPrediction {
  source: string;
  target: string;
  probability: number;
  time?: number;
}

interface TSGEOutputParams {
  node_embeddings: NodeEmbedding[];
  link_predictions?: LinkPrediction[];
  algorithm: string;
  parameters: Record<string, any>;
  train_loss_history?: number[];
}

interface InputParams {
  tsge_params: TSGEInputParams;
}

interface OutputParams {
  tsge_results: TSGEOutputParams;
}

interface AlgorithmRequest {
  task_id: string;
  task_callback_url?: string;
  input_params: InputParams;
}

interface AlgorithmResponse {
  task_id: string;
  task_callback_url?: string;
  task_status: TaskStatus;
  task_progress: number;
  task_logs?: string;
  error_message?: string;
  output_params: OutputParams;
  metrics: any[]; // 这里可以根据Chart类型进一步定义
}

// 修改算法实现
export const TemporalSpatialGraphEmbedding: Algorithm = {
  name: "时序空间图嵌入算法(TSGE)",
  description: "基于时间序列的异构图表示学习算法，用于学习图中节点的向量表示并进行链接预测",
  configOptions: [
    {
      id: "embedding_dim",
      label: "嵌入维度",
      type: "number", 
      defaultValue: 128,
      placeholder: "128"
    },
    {
        id: "node_types",
        label: "节点类型信息",
        type: "textarea",
        placeholder: `{
        "1": "user",
        "2": "item",
        "3": "user",
        "4": "item"
      }`
    },
    {
      id: "learning_rate",
      label: "学习率",
      type: "number",
      defaultValue: 0.01,
      placeholder: "0.01"
    },
    {
      id: "epochs",
      label: "训练轮数",
      type: "number", 
      defaultValue: 100,
      placeholder: "100"
    },
    {
      id: "inputData",
      label: "输入数据",
      type: "textarea",
      placeholder: `请输入JSON格式的时序图数据序列，例如：
[
  {
    "nodes": [
      {
        "id": "n1",
        "type": "用户",
        "attributes": {"feature1": 0.8}
      }
    ],
    "edges": [
      {
        "source": "n1",
        "target": "n2",
        "type": "关系",
        "weight": 1.0
      }
    ]
  }
]`
    }
  ],

  execute: async (config) => {
    try {
      const taskId = `tsge_${Date.now()}_${Math.random().toString(36).substr(2, 9)}`
      
      let graphSequence
      try {
        graphSequence = JSON.parse(config.inputData)
      } catch (e) {
        throw new Error("输入数据必须是有效的JSON格式")
      }

      const requestData: AlgorithmRequest = {
        task_id: taskId,
        task_callback_url: "http://172.16.4.151:1122/api/tsge/callback",
        input_params: {
          tsge_params: {
            graph_sequence: graphSequence,
            embedding_dim: Number(config.embedding_dim) || 128,
            learning_rate: Number(config.learning_rate) || 0.01,
            epochs: Number(config.epochs) || 100
          }
        }
      }

      // 创建任务
      const createResponse = await fetch("http://172.16.4.151:1122/api/tsge/create", {
        method: "POST",
        headers: {
          "Content-Type": "application/json",
        },
        body: JSON.stringify(requestData)
      })

      if (!createResponse.ok) {
        const errorData = await createResponse.json()
        throw new Error(errorData.detail || "创建任务失败")
      }

      // 执行任务
      const executeResponse = await fetch(`http://172.16.4.151:1122/api/tsge/execute/${taskId}`, {
        method: "POST"
      })

      if (!executeResponse.ok) {
        const errorData = await executeResponse.json()
        throw new Error(errorData.detail || "执行任务失败")
      }

      // 轮询获取结果
      let result
      let attempts = 0
      const maxAttempts = 60

      while (attempts < maxAttempts) {
        const getResultResponse = await fetch(`http://172.16.4.151:1122/api/tsge/result/${taskId}`)
        
        if (!getResultResponse.ok) {
          const errorData = await getResultResponse.json()
          throw new Error(errorData.detail || "获取结果失败")
        }

        result = await getResultResponse.json()
        
        if (result.task_status === "COMPLETED") {
          break
        } else if (result.task_status === "FAILED") {
          throw new Error(result.error_message || "任务执行失败")
        }

        await new Promise(resolve => setTimeout(resolve, 1000))
        attempts++
      }

      if (attempts >= maxAttempts) {
        throw new Error("任务执行超时")
      }

      const outputParams = result.output_params.tsge_results
      
      // 按照时间点对链接预测结果进行分组
      const groupedPredictions = outputParams.link_predictions?.reduce((acc: Record<string, LinkPrediction[]>, pred: LinkPrediction) => {
        const key = pred.time?.toString() || 'default'
        if (!acc[key]) {
          acc[key] = []
        }
        acc[key].push(pred)
        return acc
      }, {})

      // 生成链接预测结果文本
      const linkPredictionText = groupedPredictions 
        ? Object.entries(groupedPredictions)
          .map(([timePoint, predictions]) => {
            // 添加显式类型断言
            const preds = predictions as LinkPrediction[];
            const topPredictions = preds
              .sort((a, b) => b.probability - a.probability)
              .slice(0, 5)
              .map((p) => 
                `    ${p.source} -> ${p.target} (概率: ${(p.probability * 100).toFixed(2)}%)`
              )
              .join('\n')
            
            return `  时间点 ${timePoint}:\n${topPredictions}`
          })
          .join('\n\n')
        : '无链接预测结果'

      // 生成结果摘要
      const resultText = `时序图嵌入学习完成：
节点嵌入数量: ${outputParams.node_embeddings.length}
嵌入维度: ${outputParams.parameters.embedding_dim}
算法: ${outputParams.algorithm}
训练参数:
  - 学习率: ${outputParams.parameters.learning_rate}
  - 训练轮数: ${outputParams.parameters.epochs}

链接预测结果:
${linkPredictionText}`

      return {
        result: resultText,
        steps: [],
        executionTime: 0,
        details: {
          node_embeddings: outputParams.node_embeddings,
          link_predictions: outputParams.link_predictions,
          parameters: outputParams.parameters,
          training_loss: outputParams.train_loss_history
        }
      }
    } catch (e: any) {
      console.error("TSGE算法执行出错:", e)
      throw e
    }
  }
} 